import ConfigSpace as CS
from math import floor, ceil
from types import SimpleNamespace


def cs_hyparams(hyp_params, search_hyp_params_list):
  # limit the ks of dyn_conv
  len_dict = {
      'l': hyp_params.l_len,
      'a': hyp_params.a_len,
      'v': hyp_params.v_len
  }
  bound_dict = {}
  for name_i, len_i in len_dict.items():
    for name_j, len_j in len_dict.items():
      # generate upper limit, consider the order first.
      ks_base = floor(len_j / len_i) // 2 * 2
      upper_bound = min(len_j - 1, 49)
      search_bound = (upper_bound - ks_base) // 2
      bound_dict[name_i + name_j] = search_bound

  # limit the width of the model
  # bug in ConfigSpace package. cannot support #heads or $dim_per_heads more than 15
  width_limit = [64, 15, 15]  # #dim, #heads, #dim_per_heads

  # genrating the init config
  cs_hyparams_config_base = {
      "lr":
          CS.UniformFloatHyperparameter("lr", lower=1e-5, upper=5e-3, log=True),
      "attn_dropout_a":
          CS.UniformFloatHyperparameter("attn_dropout_a", lower=0.0, upper=0.3),
      "attn_dropout_v":
          CS.UniformFloatHyperparameter("attn_dropout_v", lower=0.0, upper=0.3),
      "attn_dropout_l":
          CS.UniformFloatHyperparameter("attn_dropout_l", lower=0.0, upper=0.3),
      "attn_dropout":
          CS.UniformFloatHyperparameter("attn_dropout", lower=0.0, upper=0.3),
      "embed_dropout":
          CS.UniformFloatHyperparameter("embed_dropout", lower=0.0, upper=0.3),
      "out_dropout":
          CS.UniformFloatHyperparameter("out_dropout", lower=0.0, upper=0.3),
      "kernel_size_av":
          CS.UniformIntegerHyperparameter(
              "kernel_size_av", lower=1, upper=bound_dict['av']),
      "kernel_size_al":
          CS.UniformIntegerHyperparameter(
              "kernel_size_al", lower=1, upper=bound_dict['al']),
      "kernel_size_la":
          CS.UniformIntegerHyperparameter(
              "kernel_size_la", lower=1, upper=bound_dict['la']),
      "kernel_size_lv":
          CS.UniformIntegerHyperparameter(
              "kernel_size_lv", lower=1, upper=bound_dict['lv']),
      "kernel_size_va":
          CS.UniformIntegerHyperparameter(
              "kernel_size_va", lower=1, upper=bound_dict['va']),
      "kernel_size_vl":
          CS.UniformIntegerHyperparameter(
              "kernel_size_vl", lower=1, upper=bound_dict['vl']),
      "kernel_size":
          CS.UniformIntegerHyperparameter(
              "kernel_size", lower=1, upper=min(bound_dict.values())),
      "nlayers":
          CS.UniformIntegerHyperparameter("nlayers", lower=2, upper=6),
      "num_heads":
          CS.UniformIntegerHyperparameter(
              "num_heads", lower=2, upper=width_limit[1]),
      "dims_per_head":
          CS.UniformIntegerHyperparameter(
              "dims_per_head", lower=2, upper=width_limit[2]),
  }

  cs_hyparams_list = CS.ConfigurationSpace()
  for hyp_param_name in search_hyp_params_list:
    cs_hyparams_list.add_hyperparameter(
        cs_hyparams_config_base.get(hyp_param_name))

  if 'num_heads' in search_hyp_params_list:
    # forbid num_heads * dims_per_head >64
    for i in range(2, width_limit[1] + 1):
      dim_per_heads_forbid_list = list(
          range(ceil((width_limit[0] + 1) / i), width_limit[2] + 1))
      if dim_per_heads_forbid_list != []:
        num_heads_forbidden_clause = CS.ForbiddenEqualsClause(
            cs_hyparams_config_base.get('num_heads'), i)
        dim_per_heads_forbidden_clause = CS.ForbiddenInClause(
            cs_hyparams_config_base.get('dims_per_head'),
            dim_per_heads_forbid_list)
        cs_hyparams_list.add_forbidden_clause(
            CS.ForbiddenAndConjunction(num_heads_forbidden_clause,
                                       dim_per_heads_forbidden_clause))

  return cs_hyparams_list


if __name__ == "__main__":
  hyp_params = {'l_len': 20, 'a_len': 400, 'v_len': 500}
  x = cs_hyparams(SimpleNamespace(**hyp_params), ['num_heads', 'dims_per_head'])
  while (True):
    f = x.sample_configuration()
    if f['num_heads'] * f['dims_per_head'] > 64:
      print(f)